import os
import sys
sys.path.append(".")
import cv2
import argparse
import time
import numpy as np
from multiprocessing import Pool
import re
import constant as const
import img_process
import flow
import compute_flow
import shutil

def refine_align_section(root_dir, fixed_dir, moving_dir, process_num=-1):
    print(f"fixed_dir: {fixed_dir}, moving_dir: {moving_dir}")
    # 第一步，8nm图片降采样到32nm，每个tile为1024*1024
    img_process.batch_downsample(f"{root_dir}/{fixed_dir}_8nm",
                                 f"{root_dir}/{fixed_dir}_32nm", scale=0.25, process_num=-1)
    img_process.batch_downsample(f"{root_dir}/{moving_dir}_8nm",
                                 f"{root_dir}/{moving_dir}_32nm", scale=0.25, process_num=-1)
    
    # 第二步，32nm图片降采样到256nm，每个tile为128*128
    img_process.batch_downsample(f"{root_dir}/{fixed_dir}_32nm",
                                 f"{root_dir}/{fixed_dir}_256nm", scale=0.125, process_num=-1)
    img_process.batch_downsample(f"{root_dir}/{moving_dir}_32nm",
                                 f"{root_dir}/{moving_dir}_256nm", scale=0.125, process_num=-1)
    
    # # 再将256nm tile小图合并成一张大图(这一步不需要操作)
    # img_process.merge_tiles("/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_256nm",
    #                         "/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/S11867_256nm.png")
    
    # 第三步，将每4个tile合并成1个tile, 即T(i,j), T(i+1,j), T(i,j+1), T(i+1,j+1)合并成1个新的T(i,j)，新的tile为256*256
    img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{fixed_dir}_256nm",
                                                            f"{root_dir}/{fixed_dir}_256nm_combine")
    img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{moving_dir}_256nm",
                                                            f"{root_dir}/{moving_dir}_256nm_combine")

    # 第四步，将第三步生成的两张切片的256*256小图使用深度学习框架进行配准(这一步需要GPU，生成光流场文件)
    compute_flow.process_image_pairs(f"{root_dir}/{fixed_dir}_256nm_combine",
                                     f"{root_dir}/{moving_dir}_256nm_combine",
                                     f"{root_dir}/{moving_dir}_256nm_combine_flow")

    # 第五步，根据光流场文件夹中每个flow文件(由4块小区域组成，因为第二、三步中每4个tile合并成1个tile，并且tile有重叠), 计算出每个小区域的实际光流场，再合并成一整个大的光流场
    total_flow_256nm = flow.combine_flow(f"{root_dir}/{moving_dir}_256nm_combine_flow", 128)
    # np.save("/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_256nm_flow.npy", total_flow_256nm)
    
    # 第六步，对整个光流场进行上采样到32nm，并分割为1024*1024的小光流场保存
    flow_32nm = flow.upsample_flow(total_flow_256nm, 8)  # 从256nm上采样到32nm
    flow.split_flow(flow_32nm, 1024, f"{root_dir}/{moving_dir}_32nm_coarse_flow")

    # 第七步，对上一步的小光流场进行合并，并且带overlap，同样的，对32nm图片进行相应的操作
    flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse_flow",
                                                          f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine", process_num=-1)
    img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_32nm",
                                                                 f"{root_dir}/{moving_dir}_32nm_combine", process_num=-1)
    
    # 第八步，对上一步合并的图像应用其对应的光流场得到32nm下align的粗略的结果
    flow.process_flow_to_img(f"{root_dir}/{moving_dir}_32nm_combine",
                             f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine",
                             f"{root_dir}/{moving_dir}_32nm_combine_coarse")
    # # 作为比较，可以将fixed图像也按照第七步合并成多个大图，然后对比第八步的结果
    # img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/07246_32nm",
    #                                                              f"{root_dir}/07246_32nm_combine", process_num=-1)

    # 第九步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像
    img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_32nm_combine_coarse",
                                                   f"{root_dir}/{moving_dir}_32nm_coarse",
                                                   1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)
    # # 可以合并成大图看一下32nm下的结果
    # img_process.merge_tiles("/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_32nm_coarse",
    #                         "/media/hqjin/Elements/em_data/mec/20230321_MEC/wafer52/out/align/refine/11867_32nm_coarse_img.png")

    # 第十步，将上一步的每张32nm下的小图(1024*1024)分割成16个(256*256)的小图
    img_process.split_tiles(f"{root_dir}/{moving_dir}_32nm_coarse",
                            f"{root_dir}/{moving_dir}_32nm_coarse_split", 256)
    img_process.split_tiles(f"{root_dir}/{fixed_dir}_32nm",
                            f"{root_dir}/{fixed_dir}_32nm_split", 256)
    
    # 第十一步，将上一步的每4个tile(256*256)合并成1个tile, 即T(i,j), T(i+1,j), T(i,j+1), T(i+1,j+1)合并成1个新的T(i,j)，新的tile为(512*512)
    img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse_split",
                                                            f"{root_dir}/{moving_dir}_32nm_coarse_split2")
    img_process.combine_4_tiles_to_larger_tile_with_overlap(f"{root_dir}/{fixed_dir}_32nm_split",
                                                            f"{root_dir}/{fixed_dir}_32nm_split2")

    # # 第十二步，将上一步生成的两张切片的512*512小图使用深度学习框架进行配准(这一步需要GPU，生成光流场文件)
    compute_flow.process_image_pairs(f"{root_dir}/{fixed_dir}_32nm_split2",
                                     f"{root_dir}/{moving_dir}_32nm_coarse_split2",
                                     f"{root_dir}/{moving_dir}_32nm_coarse_split2_flow")

    # 第十三步，根据光流场文件夹中每个flow文件(由4块小区域组成，因为第十、十一步中每4个tile合并成1个tile，并且tile有重叠), 计算出每个小区域的实际光流场，再合并成一整个大的光流场
    total_flow_32nm = flow.combine_flow(f"{root_dir}/{moving_dir}_32nm_coarse_split2_flow", 256)
    # 再将大的光流场分割为1024*1024的小光流场保存
    flow.split_flow(total_flow_32nm, 1024, f"{root_dir}/{moving_dir}_32nm_fine_flow")

    # 第十四步，对上一步的小光流场进行合并，并且带overlap，同样的，对32nm图片进行相应的操作
    flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_32nm_fine_flow",
                                                          f"{root_dir}/{moving_dir}_32nm_fine_flow_combine", process_num=-1)
    img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_32nm_coarse",
                                                                 f"{root_dir}/{moving_dir}_32nm_coarse_combine", process_num=-1)
    
    # 第十五步，对上一步合并的图像应用其对应的光流场得到32nm下align的粗略的结果
    flow.process_flow_to_img(f"{root_dir}/{moving_dir}_32nm_coarse_combine",
                             f"{root_dir}/{moving_dir}_32nm_fine_flow_combine",
                             f"{root_dir}/{moving_dir}_32nm_combine_fine")
    
    # 第十六步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像
    img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_32nm_combine_fine",
                                                   f"{root_dir}/{moving_dir}_32nm_fine",
                                                   1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)
    # img_process.merge_tiles(f"{root_dir}/07247_32nm_fine",
    #                         f"{root_dir}/07247_32nm_fine_img.png")

    # 第十七步，将32nm粗略的合并的光流场和精细的合并的光流场做复合运算
    flow.composite_flows(f"{root_dir}/{moving_dir}_32nm_coarse_flow_combine",
                         f"{root_dir}/{moving_dir}_32nm_fine_flow_combine",
                         f"{root_dir}/{moving_dir}_32nm_comp_flow_combine")

    # 第十八步，将32nm复合的合并的光流场拆解为一个个小的光流场(1024*1024*2)
    flow.get_small_flows_from_combine_flows(f"{root_dir}/{moving_dir}_32nm_comp_flow_combine",
                                            f"{root_dir}/{moving_dir}_32nm_comp_flow",
                                            1, 52, 1, 54, 1024, const.ROW_OVERLAP, const.COL_OVERLAP)

    # 第十九步，将上一步的32nm的一个个小光流场(1024*1024*2)进行上采样至8nm(4096*4096*2)
    flow.batch_upsample_flows(f"{root_dir}/{moving_dir}_32nm_comp_flow",
                              f"{root_dir}/{moving_dir}_8nm_flow", 4, 50)
    
    # 第二十步，将8nm下的一个个小光流场进行合并，并且带overlap，同样的，对8nm图片进行相应的操作
    flow.batch_combine_flows_to_larger_flows_with_overlap(f"{root_dir}/{moving_dir}_8nm_flow",
                                                          f"{root_dir}/{moving_dir}_8nm_flow_combine", process_num=50)
    img_process.batch_combine_tiles_to_larger_tiles_with_overlap(f"{root_dir}/{moving_dir}_8nm",
                                                                 f"{root_dir}/{moving_dir}_8nm_combine", process_num=-1)
    
    # 第二十一步，对上一步合并的图像应用其对应的光流场得到8nm下align的粗略的结果
    flow.process_flow_to_img(f"{root_dir}/{moving_dir}_8nm_combine",
                             f"{root_dir}/{moving_dir}_8nm_flow_combine",
                             f"{root_dir}/{moving_dir}_8nm_combine_fine")

    # 第二十二步，将上一步的应用光流场得到的图像还原成每一小块不重叠的图像
    img_process.get_small_tiles_from_combine_tiles(f"{root_dir}/{moving_dir}_8nm_combine_fine",
                                                   f"{root_dir}/{moving_dir}_8nm_fine",
                                                   1, 52, 1, 54, 4096, const.ROW_OVERLAP, const.COL_OVERLAP)


def main():
    parser = argparse.ArgumentParser(description="This is a auxiliary tool.")
    parser.add_argument("-r", "--root_dir", type=str, default="/CX/neuro_segment/user/jinhaiqun/out/mec/align_imgs/refine", help="source dir")
    parser.add_argument("-f", "--fixed_dir", type=str, default="01216", help="fixed image dir")
    parser.add_argument("-m", "--moving_dir", type=str, default="01217", help="moving image dir")
    parser.add_argument("-p", "--process_num", type=int, default=16, help="the number of processes to use(default: 16)")
    args = parser.parse_args()

    refine_align_section(args.root_dir, args.fixed_dir, args.moving_dir, 112)


if __name__ == '__main__':
    main()
